from pyspark.ml.linalg import Vectors
from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
from pyspark.ml import Pipeline
from pyspark.ml.feature import IndexToString, StringIndexer, VectorIndexer
from pyspark.ml.classification import DecisionTreeClassifier
spark = SparkSession.builder.master('local').appName('dicisiontree').getOrCreate()

def load_data(x):
    rel = {}
    rel['features'] = Vectors.dense(float(x[0]), float(x[1]), float(x[2]), float(x[3]))
    rel['label'] = str(x[4])
    return rel

data = spark.sparkContext.textFile('./Iris.txt').map(lambda line: line.split(',')).map(lambda p: Row(**load_data(p))).toDF()

labelIndexer=StringIndexer(inputCol='label',outputCol='indexedLabel').fit(data)
featureIndexer=VectorIndexer(inputCol='features',outputCol='indexedFeatures',maxCategories=4).fit(data)
labelConverter=IndexToString(inputCol='prediction',outputCol='predictedLabel').setLabels(labelIndexer.labels)

trainingData,testData=data.randomSplit([0.7,0.3])

dtClassifier=DecisionTreeClassifier().setLabelCol('indexedLabel').setFeaturesCol('indexedFeatures')

dtPipeline=Pipeline().setStages([labelIndexer,featureIndexer,dtClassifier,labelConverter])
dtPipelineModel=dtPipeline.fit(trainingData)
dtPredictions=dtPipelineModel.transform(testData)
dtPredictions.select('predictedLabel','label','features').show()

evaluator=MulticlassClassificationEvaluator().setLabelCol('indexedLabel').setPredictionCol('prediction')
print(evaluator.evaluate(dtPredictions))